#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <unistd.h>
#include <signal.h>
#include <sched.h>

// 全局变量
volatile sig_atomic_t stop = 0; // 用于控制线程退出
pthread_t *threads = NULL;      // 压力测试线程数组
int thread_count = 1;           // 默认线程数量

// 压力测试线程函数
void *stress_thread(void *arg) {
    int thread_id = *(int *)arg;
    char thread_name[32];
    snprintf(thread_name, sizeof(thread_name), "stress_thread%d", thread_id);
    pthread_setname_np(pthread_self(), thread_name);

    // 频繁内存读写和算术计算
    while (!stop) {
        volatile int data = 0;
        for (int i = 0; i < 1000000; i++) {
            data += i * i; // 算术计算
        }
    }

    printf("Thread %s exited.\n", thread_name);
    return NULL;
}

// 信号处理函数
void handle_signal(int signum) {
    stop = 1; // 设置退出标志
}

// 设置线程的 CPU 亲和性
void set_thread_affinity(pthread_t thread, int cpu_core) {
    // cpu_set_t cpuset;
    // CPU_ZERO(&cpuset);
    // CPU_SET(cpu_core, &cpuset);
    // if (pthread_setaffinity_np(thread, sizeof(cpuset), &cpuset) != 0) {
    //     perror("pthread_setaffinity_np");
    //     exit(EXIT_FAILURE);
    // }
}

int main(int argc, char *argv[]) {
    // 解析命令行参数
    if (argc > 1) {
        thread_count = atoi(argv[1]);
        if (thread_count <= 0) {
            fprintf(stderr, "Invalid thread count. Using default value: 1\n");
            thread_count = 1;
        }
    }

    // 注册信号处理函数
    signal(SIGINT, handle_signal);
    signal(SIGTERM, handle_signal);

    // 分配线程数组
    threads = (pthread_t *)malloc(thread_count * sizeof(pthread_t));
    if (threads == NULL) {
        perror("malloc");
        exit(EXIT_FAILURE);
    }

    // 创建压力测试线程
    for (int i = 0; i < thread_count; i++) {
        if (pthread_create(&threads[i], NULL, stress_thread, NULL) != 0) {
            perror("pthread_create");
            exit(EXIT_FAILURE);
        }
        // 设置线程的 CPU 亲和性为 0 核心
        set_thread_affinity(threads[i], 0);
    }

    // 等待所有线程结束
    for (int i = 0; i < thread_count; i++) {
        if (pthread_join(threads[i], NULL) != 0) {
            perror("pthread_join");
            exit(EXIT_FAILURE);
        }
    }

    // 释放线程数组
    free(threads);

    printf("All stress threads exited. Program terminated.\n");
    return 0;
}